
### Doubly robust estimating equation
## Input: 
##      current estimate of alpha, beta from pmle, gamma from correctly specified glm
## Output: 
##      New estimate of alpha



dr.estimate.onestep <- function(data, theta.coef, eta.coef,gop.coef, 
                                p.y.func, b.coef, 
                                cnt, rep=500, seed,
                                max.step, thres, dr.warm=NULL){
  
  N    = nrow(data)
  Ak   = cbind(data[,(pL0+1):(pL0+K+1)])
  Lk   = cbind(data[,(pL0+K+2):(pL0+2*K+2)])
  Y = data[,"Y"]
  X.base = data[,1:pL0]
  dr.objective <- function(pars){
    # pars <- theta.coef # DEBUG
    U.m <- E.U.m <- array(NA, dim=c(N, K+1))
    
    LK = Lk[,(K+1)]; AK = Ak[,(K+1)]
    U.K.cov <- cbind(1-LK, LK, X.base)
    # U.m[, K+1] <- c(Y * exp(-t(t(U.K.cov) * AK) %*% pars))
    U.m[, K+1] <- c(Y * exp(-(U.K.cov* AK) %*% pars))
    
    # A t(t()) trick for multiplying one vector to each row
    # By default d1 * v1 is multiplying v1 to each column of v1
    #test code for the t(t()) trick
    # c1 <- c(1,2,3)
    # c2 <- c(4,5,6)
    # c3 <- c(7,8,9)
    # d1 <- data.frame(c1,c2,c3)
    # v1 <- c(1,2,3)
    # t(t(d1)*v1)
    #       c1 c2 c3
    # [1,]  1  8 21
    # [2,]  2 10 24
    # [3,]  3 12 27
    # E.U.m.func <- list()
    E.U.m.func <- function(A.mminus1.bar, Lm.bar){
        Ak.K.eq.0 <- cbind(A.mminus1.bar, rep(0, N))
        return(p.y.func(A= Ak.K.eq.0, L = Lm.bar))
    }
      
    A.mminus1.bar <- Ak[,1:K]
    Lm.bar <- Lk
    E.U.m[, K+1] <- E.U.m.func(A.mminus1.bar, Lm.bar)
    E.U.mplus1.1.1 <- E.U.mplus1.0.1 <- 
      E.U.mplus1.1.0 <- E.U.mplus1.0.0 <- array(NA,dim=c(N, K+1))
    E.U.mplus1.1.1[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(1,N)) , cbind(Lk[,1:K], rep(1,N)))
    E.U.mplus1.0.1[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(1,N)) , cbind(Lk[,1:K], rep(0,N)))
    E.U.mplus1.1.0[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(0,N)) , cbind(Lk[,1:K], rep(1,N)))
    E.U.mplus1.0.0[, K+1] <- E.U.m.func(cbind(Ak[,1:(K-1)],rep(0,N)) , cbind(Lk[,1:K], rep(0,N)))  
    
    for (m in K:1){
      # m <- K #DEBUG
      # m <- 1 # DEBUG
      # print(m)
      Lm = Lk[,m]; Am = Ak[,m]
      U.m.cov <- cbind(1-Lm, Lm, X.base)
      # U.m[, m] <- U.m[, (m+1)] *  c(Y * exp(-t(t(U.m.cov) * Am) %*% pars))
      U.m[, m] <- U.m[, (m+1)] *  c(Y * exp(-(U.m.cov) * Am) %*% pars)
      
      get.m.bar <- function(df, m){
        if (m == 1){
          return(matrix(df[,1:1], ncol=1))
        }else if (m <1){
          return(NULL)
        }else{
          return(df[,1:m])
        }
      }
        A.mminus1.bar <- get.m.bar(Ak, m-1)
        A.mminus2.bar <- get.m.bar(Ak, m-2)
        L.mminus1.bar <- get.m.bar(Lk, m-1)
        Lm.bar <- get.m.bar(Lk, m)
      
      
      E.U.m.func <- function(A.mminus1.bar, Lm.bar){
        # A.mminus1.bar is of dimension c(N, m-1)
        # Lm.bar is of dimension c(N, m)
        Lm = Lm.bar[,ncol(Lm.bar)]
        p.Am <- c(prob.score.func(b.coef,L=Lm,X=X.base))
        U.m.cov <- cbind(1-Lm, Lm, X.base)
        exp.factor <- c(exp(-U.m.cov %*% theta.coef))
        
        p.Lm.1.1 <- p.L.func(eta.coef= eta.coef, ak= rep(1,N), lk= Lm, X=X.base)
        p.Lm.0.1 <- (1-p.Lm.1.1)
        p.Lm.1.0 <- p.L.func(eta.coef= eta.coef, ak= rep(0,N), lk= Lm, X=X.base)
        p.Lm.0.0 <- 1-p.Lm.1.0
        # res <- (p.Am * exp.factor * 
        #          (p.Lm.1.1*E.U.mplus1.1.1[, m+1] + p.Lm.0.1*E.U.mplus1.0.1[, m+1]) +
        #       (1- p.Am) * 
        #         (p.Lm.1.0*E.U.mplus1.1.0[, m+1] + p.Lm.0.0*E.U.mplus1.0.0[, m+1])
        #       )
        res <- (p.Lm.1.0*E.U.mplus1.1.0[, m+1] + p.Lm.0.0*E.U.mplus1.0.0[, m+1])
        return(res)
      }
      
      E.U.m[, m] <- E.U.m.func(A.mminus1.bar=A.mminus1.bar,
                                    Lm.bar = Lm.bar)
      
      E.U.mplus1.1.1[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(1,N)) , cbind(L.mminus1.bar, rep(1,N)))
      E.U.mplus1.0.1[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(1,N)) , cbind(L.mminus1.bar, rep(0,N)))
      E.U.mplus1.1.0[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(0,N)) , cbind(L.mminus1.bar, rep(1,N)))
      E.U.mplus1.0.0[, m] <- E.U.m.func(cbind(A.mminus2.bar,rep(0,N)) , cbind(L.mminus1.bar, rep(0,N)))  
    }
    
    d.m <- E.d.m <- array(NA, dim=c(N, pL0 + 2))
    total <- array(0, dim=c(N,pL0+2))
    for (m in 1:(K+1)){
      # m <- 1
      Lm <- Lk[,m]; Am <- Ak[,m]
      b.cov <- cbind(rep(1,N),Lm, X.base)
      p.Am <- c(prob.score.func(b.coef,L=Lm,X=X.base))
      
      d.m <- (b.cov)*Am
      E.d.m <- (b.cov)*p.Am
      total <- total + (d.m - E.d.m) * (U.m[,m] - E.U.m[,m])
      # d.m[,m,] <- t(t(b.cov)*Am)
      # E.d.m[,m,] <- t(t(b.cov)*p.Am)
    }
    
    tmp <- colSums(total * cnt)
    
    # for (m in 1:(K+1)){
    #   total <- total + (d.m[,m,] - E.d.m[,m,]) * (U.m[,m] - E.U.m[,m])
    # }
    
    return(sum(tmp^2))
  }
  # epsilon = 0.1
  # pars <- startpars <- runif(pL0+2,-epsilon, epsilon)
  pars <- startpars <- runif(pL0+2,-6, 10)
  if (!is.null(dr.warm)){
    pars <- startpars <- dr.warm$theta.coef[c(1:2, 4:7)]
  }
  # alpha.est <- opt$par
  Optim = function(par.update, fn){
    optim(par.update, fn,
          control=list(maxit=max.step), method =  "L-BFGS-B",
          lower = -rep(10,10), upper = rep(10,10))
  } 
  Diff = function(x,y) sum((x-y)^2)/sum(x^2+thres)
  diff = thres + 1; step = 0; total.count=0
  while(diff > thres & step < max.step){
    if(step %% 10 == 0) cat("This is the ", step, "th step. Diff = ", diff, "\n")
    step = step + 1
    opt1 = Optim(pars, dr.objective)
    diff1 = Diff(opt1$par,pars)
    pars= opt1$par
    
    diff = max(diff1)
    total.count =  total.count + opt1$counts
  }
  print(paste0("Numbers of iterations for DR is: ", total.count))
  theta.coef.dr <- pars
  
  return(theta.coef.dr)
}
